/*
 * Copyright 2023-2025 JetBrains s.r.o and contributors. Use of this source code is governed by the Apache 2.0 license.
 */

package kotlinx.rpc.krpc.test.compat

import ch.qos.logback.classic.Logger
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.ExperimentalCoroutinesApi
import kotlinx.coroutines.test.TestScope
import kotlinx.coroutines.withContext
import kotlinx.rpc.krpc.test.compat.service.TestStarter_Latest
import kotlinx.rpc.test.WaitCounter
import kotlinx.rpc.test.runTestWithCoroutinesProbes
import org.junit.jupiter.api.DynamicTest
import org.slf4j.LoggerFactory
import java.net.URLClassLoader
import java.util.stream.Stream
import kotlin.io.path.Path
import kotlin.test.assertTrue
import kotlin.time.Duration
import kotlin.time.Duration.Companion.seconds

enum class Role {
    Server, Client;
}

class VersionRolePair(
    val version: Versions,
    val role: Role,
)

@Suppress("unused")
val Versions.client get() = VersionRolePair(this, Role.Client)

@Suppress("unused")
val Versions.server get() = VersionRolePair(this, Role.Server)

abstract class KrpcProtocolCompatibilityTestsBase {
    interface LoadedStarter {
        val version: Versions
        val starter: Starter
        suspend fun close()
    }

    class LoadedStarterImpl(override val version: Versions, val classLoader: URLClassLoader) : LoadedStarter {
        override val starter = classLoader
            .loadClass("kotlinx.rpc.krpc.test.compat.service.TestStarter_${version.name}")
            .getDeclaredConstructor()
            .newInstance() as Starter

        override suspend fun close() {
            withContext(Dispatchers.IO) {
                classLoader.close()
            }
            starter.stopClient()
            starter.stopServer()
        }
    }

    private fun prepareStarters(exclude: List<Versions>): List<LoadedStarter> {
        return Versions.entries.filter { it !in exclude && it != Versions.Latest }.map { version ->
            val jarPathValue = System.getenv("JAR_PATH_${version.name}")
                ?: error("JAR_PATH_${version.name} environment variable is not set")

            val jarPath = Path(jarPathValue).toUri().toURL()

            val versionClassLoader = URLClassLoader(arrayOf(jarPath), javaClass.classLoader)

            LoadedStarterImpl(version, versionClassLoader)
        } + latestStarter()
    }

    private fun latestStarter() = object : LoadedStarter {
        override val version: Versions = Versions.Latest
        override val starter: Starter = TestStarter_Latest()

        override suspend fun close() {
            starter.stopClient()
            starter.stopServer()
        }
    }

    class TestEnv(
        val old: Starter,
        val new: Starter,
        val appender: TestLogAppender,
        val testScope: TestScope,
    ) : CoroutineScope by testScope {
        fun assertNoErrorsInLogs() {
            assertTrue { appender.errors.isEmpty() }
        }
    }

    @OptIn(ExperimentalCoroutinesApi::class)
    private fun runTest(
        role: Role,
        exclude: List<Versions>,
        timeout: Duration = 10.seconds,
        body: suspend TestEnv.() -> Unit,
    ): Stream<DynamicTest> {
        return prepareStarters(exclude).map { old ->
            DynamicTest.dynamicTest("$role ${old.version}") {
                runTestWithCoroutinesProbes(timeout = timeout) {
                    val root = LoggerFactory.getLogger(org.slf4j.Logger.ROOT_LOGGER_NAME) as Logger
                    val testAppender = root.getAppender("TEST") as TestLogAppender
                    testAppender.events.clear()
                    val new = latestStarter()
                    try {
                        val env = TestEnv(old.starter, new.starter, testAppender, this)
                        body(env)
                    } finally {
                        testAppender.events.clear()
                        old.close()
                        new.close()
                    }
                }
            }
        }.stream()
    }

    private fun testTransport() = LocalTransport()

    private fun testOldClientWithNewServer(
        perCallBufferSize: Int = 1,
        timeout: Duration = 60.seconds,
        exclude: List<Versions>,
        body: suspend TestEnv.(CompatService, CompatServiceImpl) -> Unit,
    ) = runTest(Role.Client, exclude, timeout) {
        val transport = testTransport()
        val config = TestConfig(perCallBufferSize, ::CompatWaitCounterImpl)
        val service = old.startClient(transport.client, config)
        val impl = new.startServer(transport.server, config)
        body(service, impl)
    }

    private fun testOldServersWithNewClient(
        perCallBufferSize: Int = 1,
        timeout: Duration = 60.seconds,
        exclude: List<Versions>,
        body: suspend TestEnv.(CompatService, CompatServiceImpl) -> Unit,
    ) = runTest(Role.Server, exclude, timeout) {
        val transport = testTransport()
        val config = TestConfig(perCallBufferSize, ::CompatWaitCounterImpl)
        val service = new.startClient(transport.client, config)
        val impl = old.startServer(transport.server, config)
        body(service, impl)
    }

    protected fun matrixTest(
        perCallBufferSize: Int = 1,
        timeout: Duration = 60.seconds,
        exclude: List<VersionRolePair> = emptyList(),
        body: suspend TestEnv.(CompatService, CompatServiceImpl) -> Unit,
    ): Stream<DynamicTest> {
        val clientExclude = exclude.filter { it.role == Role.Client }.map { it.version }
        val serverExclude = exclude.filter { it.role == Role.Server }.map { it.version }
        return Stream.concat(
            testOldClientWithNewServer(perCallBufferSize, timeout, clientExclude, body),
            testOldServersWithNewClient(perCallBufferSize, timeout, serverExclude, body),
        )
    }
}

class CompatWaitCounterImpl : CompatWaitCounter {
    private val counter = WaitCounter()

    override val value: Int
        get() = counter.value

    override fun increment() {
        counter.increment()
    }

    override suspend fun await(value: Int) {
        counter.await(value)
    }
}
